import numpy as np
from sklearn.metrics import accuracy_score
import pickle
from skimage.io import imread
from skimage.color import rgb2gray
import matplotlib.pyplot as plt

# 读取保存的模型
with open('sandstone_classifier.pkl', 'rb') as file:
    clf = pickle.load(file)

# 读取砂岩截面图2及其对应分区图
sandstone_image_path = 'Sandstone_2.tif'
partition_image_path = 'Sandstone_2_segment.tif'

# 读取图像
original_image = imread(sandstone_image_path, as_gray=True)
partition_image = imread(partition_image_path)

# 将彩色分区图转换为离散的类别标签
unique_colors, labels = np.unique(partition_image.reshape(-1, partition_image.shape[-1]), axis=0, return_inverse=True)

# 将砂岩截面图转换为特征矩阵
X = original_image.reshape(-1, 1)

# 使用模型进行预测
y_pred = clf.predict(X)

# 将预测结果转换回图像形状
predicted_segmentation = y_pred.reshape(original_image.shape)

# 计算准确率
accuracy = accuracy_score(labels.flatten(), y_pred)
print(f"准确率: {accuracy:.3f}")

# 显示结果
fig, axes = plt.subplots(1, 3, figsize=(15, 5))

# 原始图像
axes[0].imshow(original_image, cmap='gray')
axes[0].set_title('Original Image')
axes[0].axis('off')

# 真实分区图（灰度化以便于显示）
axes[1].imshow(rgb2gray(partition_image), cmap='gray')
axes[1].set_title('Segment (Ground Truth)')
axes[1].axis('off')

# 预测分区图（灰度化以便于显示）
axes[2].imshow(predicted_segmentation, cmap='gray')
axes[2].set_title(f'Segmentation (acc:{accuracy:.3f})')
axes[2].axis('off')

plt.show()